package com.ds.infra.rpc.core.transport.server.handler;

import com.ds.infra.rpc.core.protocol.RpcProtocol;
import com.ds.infra.rpc.core.protocol.RpcRequest;
import com.ds.infra.rpc.core.protocol.RpcResponse;
import com.google.common.util.concurrent.RateLimiter;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.handler.codec.TooLongFrameException;
import io.netty.handler.timeout.IdleStateEvent;
import io.netty.handler.timeout.IdleStateHandler;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 服务端连接管理器（含空闲检测/负载保护）
 *
 * @author WeiShaoying
 */
@Slf4j
public class RpcServerHandler extends ChannelInboundHandlerAdapter {

    private final AtomicInteger activeConnections = new AtomicInteger(0);
    private final int maxConnections;
    private final RateLimiter rateLimiter;

    public RpcServerHandler(int maxConnections, int qpsLimit) {
        this.maxConnections = maxConnections;
        this.rateLimiter = RateLimiter.create(qpsLimit);
    }

    @Override
    public void channelActive(ChannelHandlerContext ctx) {
        if (activeConnections.incrementAndGet() > maxConnections) {
            log.warn("Reject connection due to limit (active={})", activeConnections.get());
            ctx.close();
            return;
        }
        log.info("==== Client connected ({}), total connections: {}",
                ctx.channel().remoteAddress(), activeConnections.get());

        // 添加空闲检测
        ctx.pipeline().addLast(new IdleStateHandler(60, 0, 0, TimeUnit.SECONDS));
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) {
        int remaining = activeConnections.decrementAndGet();
        log.info("==== Client disconnected, remaining connections: {}", remaining);
    }

    @Override
    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) {
        if (evt instanceof IdleStateEvent) {
            log.warn("Closing idle connection: {}", ctx.channel().remoteAddress());
            ctx.close();
        }
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) {
        if (msg instanceof RpcProtocol) {
            RpcProtocol<RpcRequest> protocol = (RpcProtocol<RpcRequest>) msg;
            // 限流检查
            if (!rateLimiter.tryAcquire()) {
                sendErrorResponse(ctx, protocol, "Server busy"); // 传入完整协议对象
                return;
            }
            ctx.fireChannelRead(msg);
        }
    }

    private void sendErrorResponse(ChannelHandlerContext ctx, RpcProtocol<RpcRequest> protocol, String errorMsg) {
        RpcResponse response = new RpcResponse();
        response.setRequestId(protocol.getBody().getRequestId());
        response.setErrorMsg(errorMsg);

        // 修正协议构造方式
        RpcProtocol<RpcResponse> responseProtocol = new RpcProtocol<>(
                new RpcProtocol.Header(
                        protocol.getHeader().getMagic(),
                        protocol.getHeader().getVersion(),
                        RpcProtocol.MsgType.RESPONSE.getValue(),
                        0 // 长度由编码器填充
                ),
                response
        );
        ctx.writeAndFlush(responseProtocol);
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        if (cause instanceof TooLongFrameException) {
            log.warn("Frame too long from {}", ctx.channel().remoteAddress());
        } else {
            log.error("Server handler error", cause);
        }
        ctx.close();
    }
}
